#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
# from autograd import grad
# from autograd import elementwise_grad as egrad  # for functions that vectorize over inputs
# %matplotlib inline


def fmorse(x,p=[1.0,1.0,2.0,0.0]):
    ro,De,alpha,e0 = p
    r_ = x - ro
    return De*(np.exp(-2.0*alpha*r_)-2.0*np.exp(-alpha*r_)) + e0

def sigmoid(x):
    ''' sigmoid激活函数 '''
    return 1.0/(1.0+np.exp(-x))

def sigmoid_p(x):
    ''' sigmoid函数导数函数 '''
    return sigmoid(x)*(1.0-sigmoid(x))


class NeuralNetwork:
    def __init__(self,X,Y,alpha=0.01):
        ''' 构造函数 '''
        self.X  = X
        self.Y  = Y
        self.alpha = alpha     # 学习率
        self.w1 = np.random.normal(loc=0.1, scale=0.3, size=(10,1))
        self.b1 = np.random.normal(loc=0.1, scale=0.3, size=(10,1))
        self.w2 = np.random.normal(loc=0.1, scale=0.3, size=(10, 10))
        self.b2 = np.random.normal(loc=0.1, scale=0.3, size=(10,1))
        self.w3 = np.random.normal(loc=0.1, scale=0.3, size=(1, 10))
        self.b3 = np.random.normal(loc=0.1, scale=0.3, size=(1, 1))        

    def forward(self):
        ''' 前馈计算 '''
        z1  = np.matmul(self.w1,self.X) + self.b1
        a1  = sigmoid(z1)

        z2  = np.matmul(self.w2,a1) + self.b2
        a2  = sigmoid(z2)

        z3   = np.matmul(self.w3,a2) + self.b3
        a3   = sigmoid(z3)         # 最后一层神经网络的输出
        loss = np.sum(np.square(a3 -self.Y))     # 计算损失值 

        ##### 导数计算 #####
        da3  = 2*(a3-self.Y)*sigmoid_p(z3)  # dloss/dz3

        self.dw3  = np.sum(np.matmul(da3,a2.transpose([0,2,1])),axis=0)
        self.db3  = np.sum(da3,axis=0)
        
        da2       = sigmoid_p(z2)*np.matmul(self.w3.transpose(),da3)
        self.db2  = np.sum(da2,axis=0)
        self.dw2  = np.sum(np.matmul(da2,a1.transpose([0,2,1])),axis=0)
         
        da1       = sigmoid_p(z1)*np.matmul(self.w2.transpose(),da2)
        self.db1  = np.sum(da1,axis=0)

        self.dw1  = np.sum(np.matmul(da1,X.transpose([0,2,1])),axis=0)
        self.predict = a3
        return a3,loss

    def train(self,step=10000):
        ''' 训练模块 '''
        for epoch in range(step):
            self.predict,loss = self.forward()
            self.w1 = self.w1 - self.alpha*self.dw1
            self.b1 = self.b1 - self.alpha*self.db1
            self.w2 = self.w2 - self.alpha*self.dw2
            self.b2 = self.b2 - self.alpha*self.db2
            self.w3 = self.w3 - self.alpha*self.dw3
            self.b3 = self.b3 - self.alpha*self.db3
            print('epoch {:d} {:f}'.format(epoch,loss))


X  = np.linspace(0.6,3.0,50)  # 产生一向量，做为输入数据
Y  = fmorse(X)
X_ = X.copy()
ymin = np.min(Y)
ymax = np.max(Y)
Y  = (Y-ymin)/(ymax-ymin)

# print(Y.shape)

X = np.expand_dims(X,axis=1)
X = np.expand_dims(X,axis=2)      # 扩维，方便后续计算

Y_= (fmorse(X)-ymin)/(ymax-ymin)  # 数据归一化

nn = NeuralNetwork(X,Y_)
nn.train(step=10000)

Y_predict = np.squeeze(nn.predict)

plt.figure()     
plt.plot(X_,Y,alpha=0.8,linewidth=2,linestyle='-',color='blue',label=r'$Morse$ $Function$')
plt.plot(X_,Y_predict,alpha=0.8,linewidth=2,linestyle='-',color='red',label=r'$Neural$ $Network$')
plt.legend(loc='best',edgecolor='yellowgreen')
plt.show() # plt.savefig('Morse.png')
plt.close()

